import serial
import time
import threading
from  time import sleep

class RoboticArmController:
    def __init__(self, port='COM4', baudrate=115200):
        """
        初始化机械臂控制器
        :param port: 串口号 (Windows: 'COM4')
        :param baudrate: 波特率
        """
        sleep(3)
        self.port = port
        self.baudrate = baudrate
        self.ser = None
        self.connected = False
        self.operation_in_progress = False
        self.current_operation = None

        # 定义操作序列
        self.hard_operation_commands = [
            '#001P1000T3000!',  # 关节1动作
            '#002P2200T2000!',  # 关节2动作
            '#003P2000T2000!',  # 关节3动作
            '#005P1000T2000!',  # 抓手闭合
            '#005P1500T1000!',  # 抓手微调
            '#001P1500T2000!',  # 关节1复位
            '#003P1500T2000!',  # 关节3复位
            '#002P1500T2000!',  # 关节2复位
            '#003P1000T2000!',  # 关节3准备移动
            '#000P2000T2000!',  # 基座旋转
            '#005P1200T2000!',  # 抓手松开
            # 复位序列

            '#000P1500T2000!',  # 基座复位
            '#001P1500T2000!',  # 关节1复位
            '#002P1500T2000!',  # 关节2复位
            '#003P1500T2000!',  # 关节3复位
            '#004P1500T2000!',  # 关节4复位
            '#005P1500T2000!'# 抓手复位
        ]
        self.hard_operation_delays = [0.5, 0.5, 0.5, 3, 1, 0.5, 0.5, 0.5, 0.5, 2, 1, 1, 1, 1, 1, 1]

        self.soft_operation_commands = [
            '#001P1000T3000!',  # 关节1动作
            '#002P2200T2000!',  # 关节2动作
            '#003P2000T2000!',  # 关节3动作
            '#004P2200T2000!',  # 关节4动作
            '#005P1000T2000!',  # 抓手张开
            '#002P2000T2000!',  # 关节2动作
            '#003P1800T2000!',  # 关节3动作
            '#005P1900T1000!',  # 抓手微调
            '#001P1500T2000!',  # 关节1复位
            '#003P1500T2000!',  # 关节3复位
            '#002P1500T2000!',  # 关节2复位
            '#003P1000T2000!',  # 关节3准备移动
            '#000P1000T2000!',  # 基座旋转
            '#004P1500T2000!',  # 关节4动作
            '#005P1200T2000!', # 抓手松开
            # 复位序列
            # '#000P1500T2000!',  # 基座复位
            '#001P1500T2000!',  # 关节1复位
            '#002P1500T2000!',  # 关节2复位
            '#003P1500T2000!',  # 关节3复位
            '#004P1500T2000!',  # 关节4复位
            '#005P1500T2000!', # 抓手复位
            '#000P1501T2000!'
        ]
        self.soft_operation_delays = [0.5, 0.5, 0.5, 0.5, 3, 0.5, 0.5, 1, 0.5, 0.5, 0.5, 1, 1, 1,1, 1, 1, 1, 1, 1, 1]

        # 复位序列（单独使用）
        # self.reset_commands = [
        #     '#000P1500T2000!',  # 基座复位
        #     '#001P1500T2000!',  # 关节1复位
        #     '#002P1500T2000!', # 关节2复位
        #     '#003P1500T2000!',  # 关节3复位
        #     '#004P1500T2000!',  # 关节4复位
        #     '#005P1500T2000!'  # 抓手复位
        # ]
        # self.reset_delays = [1, 1, 1, 1, 1, 1]

        # 尝试连接机械臂
        self._connect()

    def _connect(self):
        """尝试连接串口"""
        try:
            self.ser = serial.Serial(self.port, self.baudrate, timeout=1)
            time.sleep(3)  # 等待机械臂初始化
            self.connected = True
            print(f"机械臂已连接: {self.port} @ {self.baudrate} baud")
            # 连接后立即复位机械臂
            self.reset_arm()
            return True
        except serial.SerialException as e:
            print(f"机械臂连接失败: {e}")
            self.connected = False
            return False

    def _execute_operation(self, commands, delays):
        """执行操作序列"""
        if not self.connected:
            print(f"[模拟] 执行机械臂操作序列 (共 {len(commands)} 条指令)")
            for cmd, delay in zip(commands, delays):
                print(f"发送指令: {cmd} (等待 {delay}秒)")
                time.sleep(delay)
            print("模拟操作完成")
            return True

        try:
            for cmd, delay in zip(commands, delays):
                self.ser.write(cmd.encode())
                print(f"发送指令: {cmd}")
                time.sleep(delay)
            return True
        except Exception as e:
            print(f"机械臂操作出错: {e}")
            return False

    def reset_arm(self):
        """复位机械臂到初始位置"""
        if not self.connected:
            print("[模拟] 复位机械臂到初始位置")
            return True

        try:
            print("复位机械臂到初始位置...")
            for cmd, delay in zip(self.reset_commands, self.reset_delays):
                self.ser.write(cmd.encode())
                print(f"发送复位指令: {cmd}")
                time.sleep(delay)
            print("机械臂复位完成")
            return True
        except Exception as e:
            print(f"机械臂复位出错: {e}")
            return False

    def operate_for_hard_package(self):
        """执行硬包装操作"""
        if self.operation_in_progress:
            print("机械臂正在执行其他操作，请等待")
            return False

        self.operation_in_progress = True
        self.current_operation = 'hard'

        # 在新线程中执行操作
        operation_thread = threading.Thread(
            target=self._execute_operation,
            args=(self.hard_operation_commands, self.hard_operation_delays)
        )
        operation_thread.start()

        # 等待操作完成
        operation_thread.join()

        self.operation_in_progress = False
        return True

    def operate_for_soft_package(self):
        """执行软包装操作"""
        if self.operation_in_progress:
            print("机械臂正在执行其他操作，请等待")
            return False

        self.operation_in_progress = True
        self.current_operation = 'soft'

        # 在新线程中执行操作
        operation_thread = threading.Thread(
            target=self._execute_operation,
            args=(self.soft_operation_commands, self.soft_operation_delays)
        )
        operation_thread.start()

        # 等待操作完成
        operation_thread.join()

        self.operation_in_progress = False
        return True

    def operate_for_package(self, package_type):
        """根据包装类型执行操作"""
        if package_type == 'hard':
            return self.operate_for_hard_package()
        elif package_type == 'soft':
            return self.operate_for_soft_package()
        else:
            print(f"未知包装类型: {package_type}")
            return False

    def close(self):
        """关闭串口连接"""
        if self.connected and self.ser and self.ser.is_open:
            # 关闭前复位机械臂
            self.reset_arm()
            self.ser.close()
            print("串口连接已关闭")
        self.connected = False